#include "kpextMemoryResource.h"

using namespace kpext;

const vk::MemoryPropertyFlags RawDeviceMemory::DefMemProp(
        vk::MemoryPropertyFlagBits::eHostVisible |
        vk::MemoryPropertyFlagBits::eHostCoherent
        );

RawDeviceMemory::RawDeviceMemory(std::shared_ptr<vk::Device> device,
                           const vk::PhysicalDeviceMemoryProperties &memprops
                           ):
    m_supportMemProps(memprops),
    m_memReq(),
    m_memProp(),
    m_memory(nullptr),
    m_isMapped(false)
{
    m_device = device;
    if(!m_device){
        throw std::runtime_error(__FUNCTION__ " input with null pointer!");
    }
}

RawDeviceMemory::~RawDeviceMemory()
{
    if(m_memory){
        if(m_isMapped) RawDeviceMemory::unmapFromHostMemory();
        m_device->freeMemory(*m_memory);
        m_memory = nullptr;
    }
}

bool RawDeviceMemory::recreate(
        const vk::MemoryRequirements &memoryRequirements,
        vk::MemoryPropertyFlags prop)
{
    auto typeBits = memoryRequirements.memoryTypeBits;
    std::vector<uint32_t> typeIndexs;

    //  查找首个符合要求的内存堆的所有索引
    for ( uint32_t i = 0; i < m_supportMemProps.memoryTypeCount; i++ )
    {
        if ( (typeBits & 1) &&
             ((m_supportMemProps.memoryTypes[i].propertyFlags & prop) == prop )
             ){
            typeIndexs.push_back(i);
        }
    }

    if(typeIndexs.size() < 1){
        SPDLOG_ERROR("No device memory heap satisfy to recreate vulkan DeviceMemory!");
        return false;
    }

    auto memptr = std::make_shared<vk::DeviceMemory>();
    vk::Result result(vk::Result::eErrorUnknown);

    for(auto idx : typeIndexs){
        vk::MemoryAllocateInfo memoryAllocateInfo(memoryRequirements.size,
                                                  idx);
        result =
                m_device->allocateMemory(&memoryAllocateInfo,
                                         nullptr,
                                         memptr.get()
                                         );
        if(vk::Result::eSuccess == result){
            break;
        }
    }

    if(vk::Result::eSuccess != result){
        SPDLOG_ERROR("recreate vulkan DeviceMemory failed!");
        return false;
    }

    if(m_memory){
        if(m_isMapped) unmapFromHostMemory();
        m_device->freeMemory(*m_memory);
        m_memory = nullptr;
    }

    m_memory = memptr;
    m_memReq = memoryRequirements;
    m_memProp = prop;
    return true;
}

vk::DeviceMemory *RawDeviceMemory::getVKDevMem()
{
    return m_memory.get();
}

void *RawDeviceMemory::mapToHostMemory(size_t offset, size_t size)
{
    if(!m_device || !m_memory){
        return nullptr;
    }

    if(m_isMapped) unmapFromHostMemory();

    if(!(m_memProp & vk::MemoryPropertyFlagBits::eHostVisible)){
        SPDLOG_ERROR("DeviceMemory without eHostVisible, can not map to host!");
        return nullptr;
    }

    vk::DeviceSize vofs = offset, vsize = VK_WHOLE_SIZE;
    if(size != 0){
        if(offset + size > m_memReq.size) return nullptr;
        vsize = size;
    }
    if(offset >= m_memReq.size) return nullptr;

    void * ret = m_device->mapMemory(*m_memory,
                                     vofs,
                                     vsize,
                                     vk::MemoryMapFlags()
                                     );
    m_isMapped = (ret != nullptr);
    return ret;
}

void RawDeviceMemory::unmapFromHostMemory()
{
    if(!m_device || !m_memory){
        return;
    }

    if(!m_isMapped) return;

    if(!(m_memProp & vk::MemoryPropertyFlagBits::eHostVisible)){
        return;
    }

    m_device->unmapMemory(*m_memory);

    m_isMapped = false;
}

void RawDeviceMemory::flushMapped(size_t size, size_t offset)
{
    if(!m_device || !m_memory){
        return;
    }

    if(!(m_memProp & vk::MemoryPropertyFlagBits::eHostVisible)){
        return;
    }

    if(size == ~0ULL){
        size = m_memReq.size;
    }

    if(offset >= m_memReq.size){
        offset = 0;
    }

    if(offset + size > m_memReq.size){
        size = m_memReq.size - size;
    }

    vk::MappedMemoryRange range(*m_memory, offset, size);
    m_device->flushMappedMemoryRanges(1, &range);
}

void RawDeviceMemory::invalidateMapped(size_t size, size_t offset)
{
    if(!m_device || !m_memory){
        return;
    }

    if(!(m_memProp & vk::MemoryPropertyFlagBits::eHostVisible)){
        return;
    }

    if(size == ~0ULL){
        size = m_memReq.size;
    }

    if(offset >= m_memReq.size){
        offset = 0;
    }

    if(offset + size > m_memReq.size){
        size = m_memReq.size - size;
    }

    vk::MappedMemoryRange range(*m_memory, offset, size);
    m_device->invalidateMappedMemoryRanges(1, &range);
}
